#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""目标检测引擎模块"""

import cv2
import numpy as np
import logging
import time
from threading import Lock
from typing import List, Dict, Tuple, Optional, Any
from datetime import datetime

# 设置日志记录器
logger = logging.getLogger(__name__)

class DetectionEngine:
    """目标检测引擎，负责加载YOLO模型并执行目标检测"""
    
    def __init__(self):
        """初始化检测引擎"""
        logger.info("开始初始化检测引擎...")
        self.lock = Lock()
        self.model = None
        self.is_initialized = False
        self.last_detection_time = 0
        self.detection_count = 0
        self.fps = 0.0
        
        # 尝试导入YOLO模型
        try:
            self._load_model()
            self.is_initialized = True
            logger.info("检测引擎初始化成功")
        except Exception as e:
            logger.error(f"检测引擎初始化失败: {str(e)}")
            self.is_initialized = False
    
    def _load_model(self):
        """加载YOLO模型
        
        尝试导入ultralytics的YOLO模型，如果失败则使用模拟模型
        """
        try:
            # 尝试导入ultralytics
            from ultralytics import YOLO
            
            # 这里应该加载实际的模型文件
            # 由于是示例，我们暂时使用一个模拟模型
            logger.info("正在加载YOLO模型...")
            
            # 这里只是示例，实际使用时需要替换为真实的模型路径
            # 比如: self.model = YOLO('path/to/helmet_detection_model.pt')
            
            # 使用模拟模型作为后备
            self.model = MockDetectionModel()
            logger.info("模拟检测模型已加载")
            
        except ImportError as e:
            logger.warning("未找到ultralytics库，使用模拟检测模型")
            self.model = MockDetectionModel()
        except Exception as e:
            logger.error(f"加载模型失败: {str(e)}")
            self.model = MockDetectionModel()
    
    def detect(self, frame: np.ndarray, **kwargs) -> List[Dict]:
        """执行目标检测
        
        Args:
            frame: 输入视频帧
            **kwargs: 额外参数
                - confidence_threshold: 置信度阈值
                - classes: 要检测的类别列表
                - agnostic_nms: 是否使用类别无关的NMS
        
        Returns:
            检测结果列表，每个结果包含bbox、置信度和类别信息
        """
        if not self.is_initialized or self.model is None:
            logger.error("检测引擎未初始化，无法执行检测")
            return []
        
        start_time = time.time()
        
        try:
            # 验证输入帧
            if frame is None or not isinstance(frame, np.ndarray) or frame.size == 0:
                logger.error("无效的输入帧")
                return []
            
            # 设置检测参数
            confidence_threshold = kwargs.get('confidence_threshold', 0.25)
            classes = kwargs.get('classes', None)
            agnostic_nms = kwargs.get('agnostic_nms', False)
            
            with self.lock:
                # 执行检测
                results = self.model.detect(
                    frame, 
                    confidence_threshold=confidence_threshold,
                    classes=classes,
                    agnostic_nms=agnostic_nms
                )
            
            # 更新统计信息
            self.detection_count += 1
            self.last_detection_time = time.time()
            self._update_fps(start_time)
            
            return results
            
        except Exception as e:
            logger.error(f"执行检测失败: {str(e)}")
            return []
    
    def _update_fps(self, start_time: float):
        """更新FPS统计
        
        Args:
            start_time: 检测开始时间
        """
        elapsed = time.time() - start_time
        if elapsed > 0:
            current_fps = 1.0 / elapsed
            # 使用移动平均平滑FPS
            self.fps = self.fps * 0.9 + current_fps * 0.1
    
    def get_status(self) -> Dict[str, Any]:
        """获取检测引擎状态
        
        Returns:
            包含引擎状态信息的字典
        """
        with self.lock:
            status = {
                'initialized': self.is_initialized,
                'model_type': 'Mock YOLO Model' if isinstance(self.model, MockDetectionModel) else 'YOLO Model',
                'detection_count': self.detection_count,
                'fps': round(self.fps, 2),
                'last_detection_time': datetime.fromtimestamp(self.last_detection_time).strftime('%Y-%m-%d %H:%M:%S') if self.last_detection_time > 0 else 'Never'
            }
        return status
    
    def release(self):
        """释放检测引擎资源"""
        try:
            with self.lock:
                if self.model is not None:
                    if hasattr(self.model, 'release'):
                        self.model.release()
                    self.model = None
            self.is_initialized = False
            logger.info("检测引擎资源已释放")
        except Exception as e:
            logger.error(f"释放检测引擎资源失败: {str(e)}")

class MockDetectionModel:
    """模拟交通场景检测模型，用于测试和开发"""
    
    def __init__(self):
        """初始化模拟模型"""
        # 交通场景相关类别
        self.class_names = ['person', 'car', 'bicycle', 'motorbike', 'truck', 'bus']
        # 为每个类别定义不同的颜色
        self.class_colors = [(0, 255, 0), (255, 0, 0), (0, 255, 255), (255, 255, 0), (255, 0, 255), (0, 128, 255)]
        
        # 为每个类别定义边界框大小范围
        self.class_box_sizes = {
            'person': {'width': (40, 80), 'height': (80, 150)},           # 人：高瘦
            'car': {'width': (80, 150), 'height': (60, 100)},             # 汽车：中等大小
            'bicycle': {'width': (60, 120), 'height': (50, 90)},          # 自行车：细长
            'motorbike': {'width': (60, 100), 'height': (50, 80)},        # 摩托车：类似自行车但稍小
            'truck': {'width': (100, 200), 'height': (70, 120)},          # 卡车：大而宽
            'bus': {'width': (120, 200), 'height': (80, 130)}             # 公交车：最大
        }
        
        # 为每个类别定义置信度范围
        self.class_confidence_ranges = {
            'person': (0.7, 0.95),
            'car': (0.75, 0.99),
            'bicycle': (0.7, 0.9),
            'motorbike': (0.7, 0.9),
            'truck': (0.8, 0.99),
            'bus': (0.8, 0.99)
        }
    
    def detect(self, frame: np.ndarray, confidence_threshold: float = 0.25, 
               classes: Optional[List[int]] = None, agnostic_nms: bool = False) -> List[Dict]:
        """模拟交通场景检测功能
        
        Args:
            frame: 输入视频帧
            confidence_threshold: 置信度阈值
            classes: 要检测的类别列表
            agnostic_nms: 是否使用类别无关的NMS
        
        Returns:
            模拟的检测结果
        """
        # 确保有一些随机性但又不会太频繁
        import random
        
        # 60%的概率检测到对象（交通场景中目标应该更常见）
        if random.random() > 0.6:
            return []
        
        results = []
        height, width = frame.shape[:2]
        
        # 生成1-4个随机检测框（交通场景中可能有更多目标）
        num_detections = random.randint(1, 4)
        
        for i in range(num_detections):
            # 随机选择类别
            class_id = random.randint(0, len(self.class_names) - 1)
            class_name = self.class_names[class_id]
            
            # 如果指定了类别列表且当前类别不在其中，则跳过
            if classes is not None and class_id not in classes:
                continue
            
            # 根据类别生成适当大小的边界框
            size_config = self.class_box_sizes[class_name]
            bbox_width = random.randint(size_config['width'][0], size_config['width'][1])
            bbox_height = random.randint(size_config['height'][0], size_config['height'][1])
            
            # 生成边界框位置，确保在画面内
            x1 = random.randint(50, width - bbox_width - 50)
            y1 = random.randint(50, height - bbox_height - 50)
            x2 = x1 + bbox_width
            y2 = y1 + bbox_height
            
            # 根据类别生成适当的置信度
            conf_range = self.class_confidence_ranges[class_name]
            confidence = random.uniform(conf_range[0], conf_range[1])
            
            # 如果置信度低于阈值，则跳过
            if confidence < confidence_threshold:
                continue
            
            # 创建检测结果
            detection = {
                'x1': x1,
                'y1': y1,
                'x2': x2,
                'y2': y2,
                'confidence': confidence,
                'class_id': class_id,
                'class_name': class_name
            }
            
            results.append(detection)
        
        return results
    
    def release(self):
        """释放模拟模型资源"""
        pass

# 创建单例实例
_detection_engine_instance = None
_detection_engine_lock = Lock()

def get_detection_engine() -> DetectionEngine:
    """获取检测引擎单例
    
    Returns:
        DetectionEngine: 检测引擎实例
    """
    global _detection_engine_instance, _detection_engine_lock
    
    with _detection_engine_lock:
        if _detection_engine_instance is None:
            _detection_engine_instance = DetectionEngine()
    
    return _detection_engine_instance